{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/ai21.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AI21\n",
    "\n",
    "This notebook shows how to use AI21's foundation models in LlamaIndex. The default model is `jamba-1.5-mini`.\n",
    "Other supported models are `jamba-1.5-large` and `jamba-instruct`. If you want to use the older Jurassic models, specify the model name `j2-mid` or `j2-ultra`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you probably need to install LlamaIndex 🦙.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-ai21"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting the AI21 API Key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When creating an `AI21` instance, you can pass the API key as a parameter. If not provided as a parameter, it defaults to the value of the environment variable `AI21_API_KEY`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from llama_index.llms.ai21 import AI21\n",
    "\n",
    "# EITHER\n",
    "api_key = <YOUR API KEY>\n",
    "os.environ[\"AI21_API_KEY\"] = api_key\n",
    "\n",
    "llm = AI21()\n",
    "\n",
    "# OR\n",
    "llm = AI21(api_key=api_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call `chat` with a list of messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Messages must be listed from oldest to newest, starting with a `user` role message and alternating between `user` and `assistant` messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.llms.ai21 import AI21\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=\"hello there\"),\n",
    "    ChatMessage(\n",
    "        role=\"assistant\", content=\"Arrrr, matey! How can I help ye today?\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name?\"),\n",
    "]\n",
    "\n",
    "# Use `preamble_override` to specify the voice and tone of the assistant.\n",
    "resp = AI21(api_key=api_key).chat(\n",
    "    messages, preamble_override=\"You are a pirate with a colorful personality\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: Arrrr, ye can call me Captain Jamba! I be a friendly pirate AI, here to help ye with any questions ye may have.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call `complete` with a prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ai21 import AI21\n",
    "\n",
    "api_key = \"Your api key\"\n",
    "resp = AI21(api_key=api_key).complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a computer scientist, entrepreneur, and writer. He is best known as the co-founder of Y Combinator, a venture capital firm that has funded over 2,000 startups, including Dropbox, Airbnb, and Reddit. Graham is also known for his essays on technology, startups, and programming languages, which he publishes on his website paulgraham.com. He is a strong advocate for the use of technology to improve people's lives and has written extensively about the importance of entrepreneurship and innovation.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Call Async Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.llms.ai21 import AI21\n",
    "\n",
    "prompt = \"What is the meaning of life?\"\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=prompt),\n",
    "]\n",
    "\n",
    "chat_resp = await AI21(api_key=api_key).achat(messages)\n",
    "\n",
    "complete_resp = await AI21(api_key=api_key).acomplete(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adjust the model behavior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Configure parameters passed to the model to adjust its behavior. For instance, setting a lower `temperature` will cause less variation between calls. Setting `temperature=0` will generate the same answer to the same question every time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ai21 import AI21\n",
    "\n",
    "llm = AI21(\n",
    "    model=\"jamba-1.5-mini\", api_key=api_key, max_tokens=100, temperature=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = llm.complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is an American computer scientist, entrepreneur, and author. He is best known for his work in the field of computer programming languages, particularly the development of the Arc programming language. He is also a co-founder of the influential startup accelerator Y Combinator, which has helped launch many successful technology startups.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stream generated responses at one token per message using the `stream_chat` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ai21 import AI21\n",
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "llm = AI21(api_key=api_key, model=\"jamba-1.5-mini\")\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "resp = llm.stream_chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None Once upon a time, in a faraway land, there was a brave and adventurous pirate named Captain Jack. He had a colorful personality and was known for his quick wit and cunning.\n",
      "\n",
      "One day, Captain Jack set sail on his trusty ship, the Black Pearl, in search of treasure. He and his crew sailed across treacherous waters and battled fierce storms, but they never gave up.\n",
      "\n",
      "After many long days at sea, they finally found the island where the treasure was said to be buried. They anchored their ship and set out on foot, armed with their trusty swords and pistols.\n",
      "\n",
      "As they made their way through the dense jungle, they encountered all manner of dangerous creatures, from venomous snakes to giant spiders. But Captain Jack and his crew were not afraid. They fought their way through, determined to reach the treasure.\n",
      "\n",
      "Finally, after what seemed like an eternity, they arrived at the spot where the treasure was supposed to be buried. They dug deep into the earth, their hearts pounding with excitement. And at last, they struck gold!\n",
      "\n",
      "They had found the treasure! Captain Jack and his crew were overjoyed. They gathered up as much gold and jewels as they could carry and set sail for home.\n",
      "\n",
      "As they sailed back to their home port, Captain Jack regaled his crew with stories of their adventures and the dangers they had overcome. They laughed and sang and drank to their good fortune.\n",
      "\n",
      "When they finally arrived back home, Captain Jack and his crew were hailed as heroes. They had risked everything to find the treasure and had returned victorious. And Captain Jack, with his colorful personality, was the most celebrated of all."
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Different models use different tokenizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ai21 import AI21\n",
    "\n",
    "llm = AI21(api_key=api_key, model=\"jamba-1.5-mini\")\n",
    "\n",
    "tokenizer = llm.tokenizer\n",
    "\n",
    "tokens = tokenizer.encode(\"Hello llama-index!\")\n",
    "\n",
    "decoded = tokenizer.decode(tokens)\n",
    "\n",
    "print(decoded)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tool Calling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent import FunctionAgent\n",
    "from llama_index.llms.ai21 import AI21\n",
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def subtract(a: int, b: int) -> int:\n",
    "    \"\"\"Subtract two integers and returns the result integer\"\"\"\n",
    "    return a - b\n",
    "\n",
    "\n",
    "def divide(a: int, b: int) -> float:\n",
    "    \"\"\"Divide two integers and returns the result float\"\"\"\n",
    "    return a - b\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "add_tool = FunctionTool.from_defaults(fn=add)\n",
    "subtract_tool = FunctionTool.from_defaults(fn=subtract)\n",
    "divide_tool = FunctionTool.from_defaults(fn=divide)\n",
    "\n",
    "llm = AI21(model=\"jamba-1.5-mini\", api_key=api_key)\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[multiply_tool, add_tool, subtract_tool, divide_tool],\n",
    "    llm=llm,\n",
    ")\n",
    "\n",
    "response = await agent.run(\n",
    "    \"My friend Moses had 10 apples. He ate 5 apples in the morning. Then he found a box with 25 apples. He divided all his apples between his 5 friends. How many apples did each friend get?\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-2x1vjWb5-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
